from langchain.chat_models import ChatOpenAI
from langchain.prompts import ChatPromptTemplate
from langchain.schema.output_parser import StrOutputParser
from langchain.schema.runnable import RunnableMap

model = ChatOpenAI()

prompt1 = ChatPromptTemplate.from_template("generate a random color")
prompt2 = ChatPromptTemplate.from_template("what is a fruit of color: {color}")
prompt3 = ChatPromptTemplate.from_template("what is countries flag that has the color: {color}")
prompt4 = ChatPromptTemplate.from_template("What is the color of {fruit} and {country}")

chain1 = prompt1 | model | StrOutputParser()
chain2 = RunnableMap({"color": chain1}) | {
    "fruit": prompt2 | model | StrOutputParser(),
    "country": prompt3 | model | StrOutputParser(),
} | prompt4

result = chain2.invoke({})
print(result)
